#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <fcntl.h>
#include <time.h>

#define TCP_PORT 5001
#define UDP_PORT 5002
#define BUFFER_SIZE 1024 * 1024

void tcp_server()
{
    int sockfd, new_socket;
    struct sockaddr_in address;
    int opt = 1;
    int addrlen = sizeof(address);
    char buffer[BUFFER_SIZE] = {0};
    int file_fd;
    ssize_t bytes_received;
    clock_t start, end;
    double cpu_time_used;

    /* 创建 TCP 套接字 */
    if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) == 0)
    {
        perror("socket failed");
        exit(EXIT_FAILURE);
    }

    /* 设置套接字选项 */
    if (setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR | SO_REUSEPORT, &opt, sizeof(opt)))
    {
        perror("setsockopt");
        exit(EXIT_FAILURE);
    }
    address.sin_family = AF_INET;
    address.sin_addr.s_addr = INADDR_ANY;
    address.sin_port = htons(TCP_PORT);

    /* 绑定地址 */
    if (bind(sockfd, (struct sockaddr *)&address, sizeof(address)) < 0)
    {
        perror("bind failed");
        exit(EXIT_FAILURE);
    }
    /* 监听连接 */
    if (listen(sockfd, 3) < 0)
    {
        perror("listen");
        exit(EXIT_FAILURE);
    }
    /* 接受连接 */
    if ((new_socket = accept(sockfd, (struct sockaddr *)&address, (socklen_t *)&addrlen)) < 0)
    {
        perror("accept");
        exit(EXIT_FAILURE);
    }

    /* 打开文件以写入接收到的数据 */
    if ((file_fd = open("received_tcp_file.bin", O_WRONLY | O_CREAT | O_TRUNC, 0644)) < 0)
    {
        perror("file open failed");
        close(sockfd);
        close(new_socket);
        exit(EXIT_FAILURE);
    }

    start = clock();
    /* 接收数据并写入文件 */
    while ((bytes_received = recv(new_socket, buffer, BUFFER_SIZE, 0)) > 0)
    {
        if (write(file_fd, buffer, bytes_received) < 0)
        {
            perror("write failed");
            break;
        }
    }
    end = clock();

    cpu_time_used = ((double)(end - start)) / CLOCKS_PER_SEC;
    printf("TCP 接收时间: %f 秒\n", cpu_time_used);

    close(file_fd);
    close(new_socket);
    close(sockfd);
}

void udp_server()
{
    int sockfd;
    struct sockaddr_in server_addr;
    char buffer[BUFFER_SIZE] = {0};
    int file_fd;
    ssize_t bytes_received;
    clock_t start, end;
    double cpu_time_used;

    /* 创建 UDP 套接字 */
    if ((sockfd = socket(AF_INET, SOCK_DGRAM, 0)) < 0)
    {
        perror("socket creation failed");
        exit(EXIT_FAILURE);
    }

    /* 绑定地址 */
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(UDP_PORT);

    if (bind(sockfd, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0)
    {
        perror("bind failed");
        close(sockfd);
        exit(EXIT_FAILURE);
    }

    /* 打开文件以写入接收到的数据 */
    if ((file_fd = open("received_udp_file.bin", O_WRONLY | O_CREAT | O_TRUNC, 0644)) < 0)
    {
        perror("file open failed");
        close(sockfd);
        exit(EXIT_FAILURE);
    }

    start = clock();
    /* 接收数据并写入文件 */
    while ((bytes_received = recvfrom(sockfd, buffer, BUFFER_SIZE, 0, NULL, NULL)) > 0)
    {
        if (write(file_fd, buffer, bytes_received) < 0)
        {
            perror("write failed");
            break;
        }
    }
    end = clock();

    cpu_time_used = ((double)(end - start)) / CLOCKS_PER_SEC;
    printf("UDP 接收时间: %f 秒\n", cpu_time_used);

    close(file_fd);
    close(sockfd);
}

int main()
{
    /* 启动 TCP 服务器 */
    tcp_server();
    /* 启动 UDP 服务器 */
    udp_server();
    return 0;
}